"""
排序操作符实现
"""

from typing import Optional, List, Dict, Any, Iterator, Tuple
import heapq
import time

from ..executor_base import IteratorExecutor
from ..types import Record, RecordBatch, ExecutionContext, SortDirection

class SortOperator(IteratorExecutor):
    """排序操作符"""
    
    def __init__(self, executor_id: str, context: ExecutionContext,
                 child: IteratorExecutor,
                 order_by: List[Dict[str, Any]]):
        super().__init__(executor_id, context)
        self.child = child
        self.order_by = order_by
        self.sorted_records: List[Record] = []
        
        # 添加子执行器
        self.add_child(child)
        
        # 验证排序规格
        self._validate_order_by()
    
    def _validate_order_by(self):
        """验证排序规格"""
        for order_spec in self.order_by:
            if 'column' not in order_spec:
                raise ValueError("Order by spec must contain 'column'")
            
            direction = order_spec.get('direction', 'ASC').upper()
            if direction not in ['ASC', 'DESC']:
                raise ValueError(f"Invalid sort direction: {direction}")
    
    def open(self):
        """打开排序操作符"""
        with self.lock:
            if self.is_open:
                return
            
            self.is_open = True
            
            # 打开子执行器
            self.child.open()
            
            # 加载并排序所有记录
            self._load_and_sort_records()
            
            # 创建排序结果迭代器
            self.current_iterator = self._create_sorted_iterator()
    
    def _load_and_sort_records(self):
        """加载并排序记录"""
        # 加载所有记录
        while True:
            child_batch = self.child.next_batch()
            if not child_batch:
                break
            
            self.sorted_records.extend(child_batch)
            self.update_stats(rows_processed=len(child_batch))
        
        # 排序记录
        sort_start_time = time.time()
        self.sorted_records.sort(key=self._create_sort_key)
        sort_time = (time.time() - sort_start_time) * 1000
        self.stats.cpu_time_ms += sort_time
        
        # 更新内存使用统计
        self.update_stats(memory_used=len(self.sorted_records) * 100)  # 估算每条记录100字节
    
    def _create_sort_key(self, record: Record) -> Tuple:
        """创建排序键"""
        key_values = []
        
        for order_spec in self.order_by:
            column = order_spec['column']
            direction = order_spec.get('direction', 'ASC').upper()
            
            value = record.get(column)
            
            # 处理None值（通常排在最后）
            if value is None:
                if direction == 'ASC':
                    value = float('inf')
                else:
                    value = float('-inf')
            
            # 处理不同数据类型
            elif isinstance(value, str):
                # 字符串比较
                if direction == 'DESC':
                    # 字符串降序需要反转比较
                    value = ''.join(chr(255 - ord(c)) for c in value)
            
            elif isinstance(value, (int, float)):
                # 数值比较
                if direction == 'DESC':
                    value = -value
            
            else:
                # 其他类型转换为字符串
                value = str(value)
                if direction == 'DESC':
                    value = ''.join(chr(255 - ord(c)) for c in value)
            
            key_values.append(value)
        
        return tuple(key_values)
    
    def _create_sorted_iterator(self) -> Iterator[Record]:
        """创建排序结果迭代器"""
        for record in self.sorted_records:
            self.check_timeout()
            self.update_stats(rows_returned=1)
            yield record
    
    def close(self):
        """关闭排序操作符"""
        with self.lock:
            if not self.is_open:
                return
            
            # 关闭子执行器
            self.child.close()
            
            # 清理排序记录
            self.sorted_records.clear()
            
            self.is_open = False
            self.current_iterator = None

class TopKSortOperator(IteratorExecutor):
    """TopK排序操作符，使用堆排序优化"""
    
    def __init__(self, executor_id: str, context: ExecutionContext,
                 child: IteratorExecutor,
                 order_by: List[Dict[str, Any]],
                 k: int):
        super().__init__(executor_id, context)
        self.child = child
        self.order_by = order_by
        self.k = k
        self.heap: List[Tuple[Tuple, Record]] = []
        
        # 添加子执行器
        self.add_child(child)
        
        # 验证排序规格
        self._validate_order_by()
    
    def _validate_order_by(self):
        """验证排序规格"""
        for order_spec in self.order_by:
            if 'column' not in order_spec:
                raise ValueError("Order by spec must contain 'column'")
    
    def open(self):
        """打开TopK排序操作符"""
        with self.lock:
            if self.is_open:
                return
            
            self.is_open = True
            
            # 打开子执行器
            self.child.open()
            
            # 构建TopK堆
            self._build_topk_heap()
            
            # 创建结果迭代器
            self.current_iterator = self._create_topk_iterator()
    
    def _build_topk_heap(self):
        """构建TopK堆"""
        while True:
            child_batch = self.child.next_batch()
            if not child_batch:
                break
            
            for record in child_batch:
                self.check_timeout()
                self.update_stats(rows_processed=1)
                
                sort_key = self._create_sort_key(record)
                
                if len(self.heap) < self.k:
                    # 堆未满，直接添加
                    heapq.heappush(self.heap, (sort_key, record))
                else:
                    # 堆已满，比较并可能替换最小元素
                    if sort_key > self.heap[0][0]:
                        heapq.heappushpop(self.heap, (sort_key, record))
    
    def _create_sort_key(self, record: Record) -> Tuple:
        """创建排序键"""
        key_values = []
        
        for order_spec in self.order_by:
            column = order_spec['column']
            direction = order_spec.get('direction', 'ASC').upper()
            
            value = record.get(column)
            
            # 处理None值
            if value is None:
                value = float('-inf') if direction == 'ASC' else float('inf')
            elif isinstance(value, (int, float)):
                # 对于最小堆，ASC需要取负值，DESC保持原值
                if direction == 'ASC':
                    value = -value
            elif isinstance(value, str):
                # 字符串处理
                if direction == 'ASC':
                    value = ''.join(chr(255 - ord(c)) for c in value)
            
            key_values.append(value)
        
        return tuple(key_values)
    
    def _create_topk_iterator(self) -> Iterator[Record]:
        """创建TopK结果迭代器"""
        # 从堆中取出所有元素并排序
        sorted_results = []
        while self.heap:
            sort_key, record = heapq.heappop(self.heap)
            sorted_results.append((sort_key, record))
        
        # 反转结果（因为使用了最小堆）
        sorted_results.reverse()
        
        # 返回排序后的记录
        for sort_key, record in sorted_results:
            self.check_timeout()
            self.update_stats(rows_returned=1)
            yield record
    
    def close(self):
        """关闭TopK排序操作符"""
        with self.lock:
            if not self.is_open:
                return
            
            # 关闭子执行器
            self.child.close()
            
            # 清理堆
            self.heap.clear()
            
            self.is_open = False
            self.current_iterator = None

class ExternalSortOperator(IteratorExecutor):
    """外部排序操作符，适用于大数据集"""
    
    def __init__(self, executor_id: str, context: ExecutionContext,
                 child: IteratorExecutor,
                 order_by: List[Dict[str, Any]],
                 memory_limit: int = 64 * 1024 * 1024):  # 64MB
        super().__init__(executor_id, context)
        self.child = child
        self.order_by = order_by
        self.memory_limit = memory_limit
        self.temp_files: List[str] = []
        self.current_memory_usage = 0
        self.chunk_size = 1000  # 每个内存块的记录数
        
        # 添加子执行器
        self.add_child(child)
    
    def open(self):
        """打开外部排序操作符"""
        with self.lock:
            if self.is_open:
                return
            
            self.is_open = True
            
            # 打开子执行器
            self.child.open()
            
            # 执行外部排序
            self._perform_external_sort()
            
            # 创建合并迭代器
            self.current_iterator = self._create_merge_iterator()
    
    def _perform_external_sort(self):
        """执行外部排序"""
        current_chunk = []
        
        while True:
            child_batch = self.child.next_batch()
            if not child_batch:
                break
            
            for record in child_batch:
                self.check_timeout()
                self.update_stats(rows_processed=1)
                
                current_chunk.append(record)
                self.current_memory_usage += self._estimate_record_size(record)
                
                # 检查是否需要写入临时文件
                if (len(current_chunk) >= self.chunk_size or 
                    self.current_memory_usage >= self.memory_limit):
                    
                    self._write_sorted_chunk(current_chunk)
                    current_chunk.clear()
                    self.current_memory_usage = 0
        
        # 处理最后一个块
        if current_chunk:
            self._write_sorted_chunk(current_chunk)
    
    def _estimate_record_size(self, record: Record) -> int:
        """估算记录大小"""
        # 简单估算：每个字段平均20字节
        return len(record) * 20
    
    def _write_sorted_chunk(self, chunk: List[Record]):
        """写入排序后的块到临时文件"""
        # 排序当前块
        chunk.sort(key=self._create_sort_key)
        
        # 这里应该写入到实际的临时文件
        # 为了简化，我们存储在内存中的列表
        temp_file_id = f"temp_{len(self.temp_files)}"
        self.temp_files.append(chunk.copy())
    
    def _create_sort_key(self, record: Record) -> Tuple:
        """创建排序键"""
        key_values = []
        
        for order_spec in self.order_by:
            column = order_spec['column']
            direction = order_spec.get('direction', 'ASC').upper()
            
            value = record.get(column)
            
            if value is None:
                value = '' if direction == 'ASC' else 'zzz'
            elif isinstance(value, (int, float)):
                if direction == 'DESC':
                    value = -value
            elif isinstance(value, str):
                if direction == 'DESC':
                    value = ''.join(chr(255 - ord(c)) for c in value)
            
            key_values.append(value)
        
        return tuple(key_values)
    
    def _create_merge_iterator(self) -> Iterator[Record]:
        """创建合并迭代器"""
        if not self.temp_files:
            return
        
        # 使用堆合并多个排序的块
        heap = []
        iterators = []
        
        # 初始化每个块的迭代器
        for i, temp_chunk in enumerate(self.temp_files):
            iterator = iter(temp_chunk)
            try:
                first_record = next(iterator)
                sort_key = self._create_sort_key(first_record)
                heapq.heappush(heap, (sort_key, i, first_record, iterator))
                iterators.append(iterator)
            except StopIteration:
                continue
        
        # 合并排序的块
        while heap:
            self.check_timeout()
            
            sort_key, chunk_id, record, iterator = heapq.heappop(heap)
            
            # 输出当前最小记录
            self.update_stats(rows_returned=1)
            yield record
            
            # 尝试从同一个块获取下一条记录
            try:
                next_record = next(iterator)
                next_sort_key = self._create_sort_key(next_record)
                heapq.heappush(heap, (next_sort_key, chunk_id, next_record, iterator))
            except StopIteration:
                # 该块已经耗尽
                continue
    
    def close(self):
        """关闭外部排序操作符"""
        with self.lock:
            if not self.is_open:
                return
            
            # 关闭子执行器
            self.child.close()
            
            # 清理临时文件
            self.temp_files.clear()
            
            self.is_open = False
            self.current_iterator = None



